{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Installation\n",
    "! pip install smolagents\n",
    "# To install from source instead of the last release, comment the command above and uncomment the following one.\n",
    "# ! pip install git+https://github.com/huggingface/smolagents.git"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Text-to-SQL"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this tutorial, we’ll see how to implement an agent that leverages SQL using `smolagents`.\n",
    "\n",
    "> Let's start with the golden question: why not keep it simple and use a standard text-to-SQL pipeline?\n",
    "\n",
    "A standard text-to-sql pipeline is brittle, since the generated SQL query can be incorrect. Even worse, the query could be incorrect, but not raise an error, instead giving some incorrect/useless outputs without raising an alarm.\n",
    "\n",
    "👉 Instead, an agent system is able to critically inspect outputs and decide if the query needs to be changed or not, thus giving it a huge performance boost.\n",
    "\n",
    "Let’s build this agent! 💪\n",
    "\n",
    "Run the line below to install required dependencies:\n",
    "```bash\n",
    "!pip install smolagents python-dotenv sqlalchemy --upgrade -q\n",
    "```\n",
    "To call Inference Providers, you will need a valid token as your environment variable `HF_TOKEN`.\n",
    "We use python-dotenv to load it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dotenv import load_dotenv\n",
    "load_dotenv()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then, we setup the SQL environment:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sqlalchemy import (\n",
    "    create_engine,\n",
    "    MetaData,\n",
    "    Table,\n",
    "    Column,\n",
    "    String,\n",
    "    Integer,\n",
    "    Float,\n",
    "    insert,\n",
    "    inspect,\n",
    "    text,\n",
    ")\n",
    "\n",
    "engine = create_engine(\"sqlite:///:memory:\")\n",
    "metadata_obj = MetaData()\n",
    "\n",
    "def insert_rows_into_table(rows, table, engine=engine):\n",
    "    for row in rows:\n",
    "        stmt = insert(table).values(**row)\n",
    "        with engine.begin() as connection:\n",
    "            connection.execute(stmt)\n",
    "\n",
    "table_name = \"receipts\"\n",
    "receipts = Table(\n",
    "    table_name,\n",
    "    metadata_obj,\n",
    "    Column(\"receipt_id\", Integer, primary_key=True),\n",
    "    Column(\"customer_name\", String(16), primary_key=True),\n",
    "    Column(\"price\", Float),\n",
    "    Column(\"tip\", Float),\n",
    ")\n",
    "metadata_obj.create_all(engine)\n",
    "\n",
    "rows = [\n",
    "    {\"receipt_id\": 1, \"customer_name\": \"Alan Payne\", \"price\": 12.06, \"tip\": 1.20},\n",
    "    {\"receipt_id\": 2, \"customer_name\": \"Alex Mason\", \"price\": 23.86, \"tip\": 0.24},\n",
    "    {\"receipt_id\": 3, \"customer_name\": \"Woodrow Wilson\", \"price\": 53.43, \"tip\": 5.43},\n",
    "    {\"receipt_id\": 4, \"customer_name\": \"Margaret James\", \"price\": 21.11, \"tip\": 1.00},\n",
    "]\n",
    "insert_rows_into_table(rows, receipts)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Build our agent"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let’s make our SQL table retrievable by a tool.\n",
    "\n",
    "The tool’s description attribute will be embedded in the LLM’s prompt by the agent system: it gives the LLM information about how to use the tool. This is where we want to describe the SQL table."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "inspector = inspect(engine)\n",
    "columns_info = [(col[\"name\"], col[\"type\"]) for col in inspector.get_columns(\"receipts\")]\n",
    "\n",
    "table_description = \"Columns:\\n\" + \"\\n\".join([f\"  - {name}: {col_type}\" for name, col_type in columns_info])\n",
    "print(table_description)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```text\n",
    "Columns:\n",
    "  - receipt_id: INTEGER\n",
    "  - customer_name: VARCHAR(16)\n",
    "  - price: FLOAT\n",
    "  - tip: FLOAT\n",
    "```\n",
    "\n",
    "Now let’s build our tool. It needs the following: (read [the tool doc](https://huggingface.co/docs/smolagents/main/en/examples/../tutorials/tools) for more detail)\n",
    "- A docstring with an `Args:` part listing arguments.\n",
    "- Type hints on both inputs and output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from smolagents import tool\n",
    "\n",
    "@tool\n",
    "def sql_engine(query: str) -> str:\n",
    "    \"\"\"\n",
    "    Allows you to perform SQL queries on the table. Returns a string representation of the result.\n",
    "    The table is named 'receipts'. Its description is as follows:\n",
    "        Columns:\n",
    "        - receipt_id: INTEGER\n",
    "        - customer_name: VARCHAR(16)\n",
    "        - price: FLOAT\n",
    "        - tip: FLOAT\n",
    "\n",
    "    Args:\n",
    "        query: The query to perform. This should be correct SQL.\n",
    "    \"\"\"\n",
    "    output = \"\"\n",
    "    with engine.connect() as con:\n",
    "        rows = con.execute(text(query))\n",
    "        for row in rows:\n",
    "            output += \"\\n\" + str(row)\n",
    "    return output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let us create an agent that leverages this tool.\n",
    "\n",
    "We use the `CodeAgent`, which is smolagents’ main agent class: an agent that writes actions in code and can iterate on previous output according to the ReAct framework.\n",
    "\n",
    "The model is the LLM that powers the agent system. `InferenceClientModel` allows you to call LLMs using HF’s Inference API, either via Serverless or Dedicated endpoint, but you could also use any proprietary API."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from smolagents import CodeAgent, InferenceClientModel\n",
    "\n",
    "agent = CodeAgent(\n",
    "    tools=[sql_engine],\n",
    "    model=InferenceClientModel(model_id=\"meta-llama/Llama-3.1-8B-Instruct\"),\n",
    ")\n",
    "agent.run(\"Can you give me the name of the client who got the most expensive receipt?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Level 2: Table joins"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let’s make it more challenging! We want our agent to handle joins across multiple tables.\n",
    "\n",
    "So let’s make a second table recording the names of waiters for each receipt_id!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "table_name = \"waiters\"\n",
    "waiters = Table(\n",
    "    table_name,\n",
    "    metadata_obj,\n",
    "    Column(\"receipt_id\", Integer, primary_key=True),\n",
    "    Column(\"waiter_name\", String(16), primary_key=True),\n",
    ")\n",
    "metadata_obj.create_all(engine)\n",
    "\n",
    "rows = [\n",
    "    {\"receipt_id\": 1, \"waiter_name\": \"Corey Johnson\"},\n",
    "    {\"receipt_id\": 2, \"waiter_name\": \"Michael Watts\"},\n",
    "    {\"receipt_id\": 3, \"waiter_name\": \"Michael Watts\"},\n",
    "    {\"receipt_id\": 4, \"waiter_name\": \"Margaret James\"},\n",
    "]\n",
    "insert_rows_into_table(rows, waiters)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since we changed the table, we update the `SQLExecutorTool` with this table’s description to let the LLM properly leverage information from this table."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "updated_description = \"\"\"Allows you to perform SQL queries on the table. Beware that this tool's output is a string representation of the execution output.\n",
    "It can use the following tables:\"\"\"\n",
    "\n",
    "inspector = inspect(engine)\n",
    "for table in [\"receipts\", \"waiters\"]:\n",
    "    columns_info = [(col[\"name\"], col[\"type\"]) for col in inspector.get_columns(table)]\n",
    "\n",
    "    table_description = f\"Table '{table}':\\n\"\n",
    "\n",
    "    table_description += \"Columns:\\n\" + \"\\n\".join([f\"  - {name}: {col_type}\" for name, col_type in columns_info])\n",
    "    updated_description += \"\\n\\n\" + table_description\n",
    "\n",
    "print(updated_description)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since this request is a bit harder than the previous one, we’ll switch the LLM engine to use the more powerful [Qwen/Qwen3-Next-80B-A3B-Thinking](https://huggingface.co/Qwen/Qwen3-Next-80B-A3B-Thinking)!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sql_engine.description = updated_description\n",
    "\n",
    "agent = CodeAgent(\n",
    "    tools=[sql_engine],\n",
    "    model=InferenceClientModel(model_id=\"Qwen/Qwen3-Next-80B-A3B-Thinking\"),\n",
    ")\n",
    "\n",
    "agent.run(\"Which waiter got more total money from tips?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It directly works! The setup was surprisingly simple, wasn’t it?\n",
    "\n",
    "This example is done! We've touched upon these concepts:\n",
    "- Building new tools.\n",
    "- Updating a tool's description.\n",
    "- Switching to a stronger LLM helps agent reasoning.\n",
    "\n",
    "✅ Now you can go build this text-to-SQL system you’ve always dreamt of! ✨"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 4
}
